"""
Spatial maps of differences between 
HYCOM & HRET during SWOT SCIENCE orbit

B. Yadidya
Feb 24, 2026
"""

import sys
sys.path.append('.')
from common_imports import *

dir = "data/"
ssh_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SLOPE_VR_v2.0.1.nc')

sss_diff_x = (sss_ds.fsss_var_x_hycom_coh5td - sss_ds.fsss_var_x_hret)
sss_diff_y = (sss_ds.fsss_var_y_hycom_coh5td - sss_ds.fsss_var_y_hret)
ssh_diff = (ssh_ds.fssh_var_hycom_coh5td - ssh_ds.fssh_var_hret)

"""
# I used the 4km regridded datasets for the plots shown in the manuscript. 
# However for storage regions, 0.25 deg regridded datasets are uploaded on Dataverse.
ssh_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SLOPE_VR_v2.0.1.nc')

sss_diff_x = coarsen_data(sss_ds.fsss_var_x_hycom_coh5td - sss_ds.fsss_var_x_hret, .1)
sss_diff_y = coarsen_data(sss_ds.fsss_var_y_hycom_coh5td - sss_ds.fsss_var_y_hret, .1)
ssh_diff = coarsen_data(ssh_ds.fssh_var_hycom_coh5td - ssh_ds.fssh_var_hret, .1)
"""

fig_width_cm = 18.4
fig_width_in = fig_width_cm / 2.54
fig_height_max_cm = 19.0
fig_height_in = fig_height_max_cm / 2.54 

fig, axs = plt.subplots(
    3, 1, 
    figsize=(fig_width_in, fig_height_in), 
    constrained_layout=True
)
axs = axs.ravel()

# --- 4. Plotting ---
levels1 = np.linspace(-1, 1, 21)
plot1 = ssh_diff.plot.contourf(ax=axs[0], levels=levels1, 
                         x='longitude', cmap=colormaps.prinsenvlag_r,
                         cbar_kwargs={'shrink': .7, 'pad': .01})

# Update colorbar to 8pt and ensure line weights
plot1.colorbar.set_label('cm$^2$', size=8)
plot1.colorbar.set_ticks([-1, 0, 1])
plot1.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0) 
plot1.colorbar.ax.minorticks_off()
plot1.colorbar.outline.set_linewidth(0.5)

limit = 2
plot2 = (sss_diff_x/1e-9).plot.contourf(ax=axs[1], levels=np.linspace(-limit, limit, 21),
                               x='longitude', cmap=colormaps.prinsenvlag_r,
                               cbar_kwargs={'shrink': .7, 'pad': .01})

plot3 = (sss_diff_y/1e-9).plot.contourf(ax=axs[2], levels=np.linspace(-limit, limit, 21),
                               x='longitude', cmap=colormaps.prinsenvlag_r,
                               cbar_kwargs={'shrink': .7, 'pad': .01})     

# Update colorbar 2
plot2.colorbar.set_ticks([-limit, 0, limit])
plot2.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0)
plot2.colorbar.ax.minorticks_off()
plot2.colorbar.set_label('10$^{-3}$(cm/km)$^2$', size=8)
plot2.colorbar.outline.set_linewidth(0.5)

# Update colorbar 3
plot3.colorbar.set_label('10$^{-3}$(cm/km)$^2$', size=8)
plot3.colorbar.set_ticks([-limit, 0, limit])
plot3.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0)
plot3.colorbar.ax.minorticks_off()
plot3.colorbar.outline.set_linewidth(0.5)


# --- 5. Final Formatting ---
axs[0].set_title('Explained SSH variance (HYCOM - HRET)', fontsize=9)
axs[1].set_title('Explained cross-track slope variance (HYCOM - HRET)', fontsize=9)
axs[2].set_title('Explained along-track slope variance (HYCOM - HRET)', fontsize=9)    

for ax in axs:
    plot_global_map_on_axes_basemap(ax) 
    ax.set_aspect('equal')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_ylim(-60, 60)

# --- 6. Panel Labels with Specific Bold Font ---
panel_label_props = dict(
    fontproperties=bold_font_props,  # <--- Uses the loaded Bold OTF file
    fontsize=9,
    ha='center',
    va='center',
    bbox=dict(
        boxstyle='round, pad=0.1',
        facecolor='white',
        edgecolor='none',
        alpha=1
    )
)

for i, label in enumerate(['A', 'B', 'C']):
    axs[i].text(0.017, .94, label, transform=axs[i].transAxes, **panel_label_props)

plt.savefig('Fig2_HYCOM_HRET_difference.png', dpi=500, bbox_inches='tight')
